from typing import TypedDict
from langgraph.graph import StateGraph
from src.common.logger import getLogger
from langgraph.constants import START, END
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

logger = getLogger()

class CalculateState(TypedDict):
    query: str
    expression: str
    answer: str

class CalculateAgent:

    def __init__(self, llm_model, tool):
        self.llm_model = llm_model
        self.tool = tool

    def extract_node(self, state: CalculateState):
        logger.info("CalculateAgent extract_node start")
        query = state["query"]
        extract_template = """
            你是一位数学计算高手，根据用户的输入问题提取其中的数学计算表达式。
            用户输入的问题：{question}
        """
        extract_prompt = ChatPromptTemplate.from_template(extract_template)
        extract_chain = extract_prompt | self.llm_model | StrOutputParser()
        extract_result = extract_chain.invoke({ "question": query })
        logger.info(f"CalculateAgent extract_node extract_result: {extract_result}")
        return { "expression": extract_result }

    def calculate_node(self, state: CalculateState):
        logger.info("CalculateAgent calculate_node start")
        expression = state["expression"]
        result = self.tool.func(expression)
        logger.info(f"CalculateAgent calculate_node result: {result}")
        return { "answer": result }

    def build_graph(self):
        logger.info("CalculateAgent build_graph start")
        graph = StateGraph(CalculateState)
        graph.add_node("extract", self.extract_node)
        graph.add_node("calculate", self.calculate_node)

        graph.add_edge(START, "extract")
        graph.add_edge("extract", "calculate")
        graph.add_edge("calculate", END)
        return graph.compile()

    def invoke(self, query):
        logger.info(f"CalculateAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        logger.info(f"CalculateAgent invoke response: {response}")
        return { "document": response.get("expression", None), "answer": response.get("answer", None) }
